﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using System.Diagnostics.Contracts;

using SimpleIR;
using PInfoModel;

namespace PIPostAnalysis.StructureSpecs
{
	/// <summary>
	/// Code foe constructing/extending specifications based on a single heap model and for computing the universal 
	/// specs for method entries.
	/// </summary>
	public static class StructSpecExtractor
	{
		public static StructureSpec ProcessProgram(CompleteAnalysisInfoForProgram pip)
		{
			StructureSpec spec = new StructureSpec();

			foreach(CompleteAnalysisInfoMethod cam in pip.MethodInfo.Values)
			{
				if(cam.InvokeID.DefInClass.TypeDescription.IsBuiltinType)
					continue;

				//if((pip.Program.MethodBodyTable.GetBodyForInvokeExactSigMatch(cam.InvokeID).Qualifier & MethodQualifiers.PrivateTag) != MethodQualifiers.Clear)
				//	continue;

				if(cam.InvokeID.Sig.IsConstructor || cam.InvokeID.Sig.InvokeName.StartsWith("get_") || cam.InvokeID.Sig.InvokeName.StartsWith("set_"))
					continue;

				foreach(CompleteAnalysisInfoMethodContext cim in cam.AllCallContexts)
					spec.ProcessSinglePIModel(cim.InState);
			}

			return spec;
		}

		public static Dictionary<PIMNode, NodeStructureSpec> ComputeStructSpec(PIModel pim)
		{
			Dictionary<PIMNode, HashSet<PIMNode>> dominfo = ComputeDominators(pim);

			Dictionary<PIMNode, NodeStructureSpec> nssl = new Dictionary<PIMNode,NodeStructureSpec>();
			InitializeExposureInfo(pim, nssl, dominfo);
			PropagateExposureInfo(pim, nssl, dominfo);

			return nssl;
		}

		/// <summary>
		/// Initialize the exposed info.
		/// </summary>
		private static void InitializeExposureInfo(PIModel pim, Dictionary<PIMNode, NodeStructureSpec> structinfo, Dictionary<PIMNode, HashSet<PIMNode>> dominfo)
		{
			foreach(PIMNode nn in pim.Nodes)
			{
				if(nn.HomeLocFor != null)
					continue;

				NodeStructureSpec nss = new NodeStructureSpec(nn);
				structinfo.Add(nn, nss);

				foreach(PIMEdge oe in nn.OutEdges)
				{
					PIMNode tn = pim.Nodes[oe.TargetAddress];

					if(!oe.IsInjective || tn.InEdges.Count((PIMEdge ee) => !ee.IsEvalStackSource && pim.Nodes[ee.SourceAddress].HomeLocFor == null) != 1)
						nss.AddExposed(oe.OffsetInSource, tn);
				}
			}
		}

		/// <summary>
		/// Take the initial exposed info and propagate the implications.
		/// </summary>
		private static void PropagateExposureInfo(PIModel pim, Dictionary<PIMNode, NodeStructureSpec> structinfo, Dictionary<PIMNode, HashSet<PIMNode>> dominfo)
		{
			Queue<PIMNode> wlist = new Queue<PIMNode>();
			foreach(PIMNode nn in pim.Nodes)
			{
				if(nn.HomeLocFor == null)
					wlist.Enqueue(nn);
			}

			while(wlist.Count != 0)
			{
				PIMNode n = wlist.Dequeue();
				bool changed = PropagateExposureInfoSingleNode(pim, n, structinfo, dominfo);

				if(changed)
				{
					foreach(PIMEdge ee in n.InEdges)
					{
						if(ee.IsEvalStackSource)
							continue;

						PIMNode snode = pim.Nodes[ee.SourceAddress];
						if(snode != n && snode.HomeLocFor == null && !wlist.Contains(snode))
							wlist.Enqueue(snode);
					}
				}
			}
		}

		/// <summary>
		/// Propagate the exposed info for a single node.
		/// </summary>
		private static bool PropagateExposureInfoSingleNode(PIModel pim, PIMNode nn, Dictionary<PIMNode, NodeStructureSpec> structinfo, Dictionary<PIMNode, HashSet<PIMNode>> dominfo)
		{
			NodeStructureSpec nnspec = structinfo[nn];

			bool changed = false;
			foreach(PIMEdge ee in nn.OutEdges)
			{
				PIMNode tn = pim.Nodes[ee.TargetAddress];
				NodeStructureSpec tnspec = structinfo[tn];

				changed |= nnspec.PropagateSpecsInto(tnspec, ee, pim, dominfo);
			}

			return changed;
		}

		/// <summary>
		/// Compute the set of nodes that a given node dominates.
		/// </summary>
		private static Dictionary<PIMNode, HashSet<PIMNode>> ComputeDominators(PIModel pim)
		{
			Dictionary<PIMNode, HashSet<PIMNode>> preddict = new Dictionary<PIMNode, HashSet<PIMNode>>();
			foreach(PIMNode n in pim.Nodes)
			{
				if(n.HomeLocFor == null)
					preddict.Add(n, new HashSet<PIMNode>());
			}

			foreach(PIMNode n in pim.Nodes)
			{
				if(n.HomeLocFor != null)
					continue;

				foreach(PIMEdge trgte in n.OutEdges)
				{
					PIMNode trgtn = pim.Nodes[trgte.TargetAddress];
					preddict[trgtn].Add(n);
				}
			}

			return ComputeAllNodesDomByDict(pim, preddict);
		}

		/// <summary>
		/// Compute a dictionary that maps from PIMNodes to the set of PIMNodes that it immediately dominataes (and self).
		/// </summary>
		private static Dictionary<PIMNode, HashSet<PIMNode>> ComputeAllNodesDomByDict(PIModel pim, Dictionary<PIMNode, HashSet<PIMNode>> preddict)
		{
			Dictionary<PIMNode, HashSet<PIMNode>> doms = ComputeDominatorSets(pim, preddict);

			Dictionary<PIMNode, HashSet<PIMNode>> resd = new Dictionary<PIMNode, HashSet<PIMNode>>();
			foreach(PIMNode nn in pim.Nodes)
			{
				HashSet<PIMNode> nns = new HashSet<PIMNode>();
				nns.Add(nn);
				resd.Add(nn, nns);
			}

			foreach(KeyValuePair<PIMNode, HashSet<PIMNode>> domsp in doms)
			{
				foreach(PIMNode nn in domsp.Value)
					resd[nn].Add(domsp.Key);
			}

			return resd;
		}

		/// <summary>
		/// Compute the dominator sets for each node.
		/// </summary>
		private static Dictionary<PIMNode, HashSet<PIMNode>> ComputeDominatorSets(PIModel pim, Dictionary<PIMNode, HashSet<PIMNode>> preddict)
		{
			//
			//Impl. taken from Advanced Compiler Design and Implelentation p. 182
			//
			Dictionary<PIMNode, HashSet<PIMNode>> domin = new Dictionary<PIMNode, HashSet<PIMNode>>();
			HashSet<PIMNode> N = new HashSet<PIMNode>();
			foreach(PIMNode nn in pim.Nodes)
			{
				if(nn.HomeLocFor == null)
					N.Add(nn);
			}

			foreach(PIMNode nn in N)
			{
				HashSet<PIMNode> doms = new HashSet<PIMNode>();
				if(nn.InEdges.Count((PIMEdge ee) => pim.Nodes[ee.SourceAddress].HomeLocFor == null) == 0)
					doms.Add(nn);
				else
					doms.UnionWith(N);
				domin.Add(nn, doms);
			}

			bool change = true;
			while(change)
			{
				change = false;
				foreach(PIMNode nn in N)
				{
					if(nn.InEdges.Count((PIMEdge ee) => pim.Nodes[ee.SourceAddress].HomeLocFor == null) == 0)
						continue;

					HashSet<PIMNode> T = new HashSet<PIMNode>(N);
					foreach(PIMNode predn in preddict[nn])
						T.IntersectWith(domin[predn]);

					HashSet<PIMNode> D = T;
					D.Add(nn);
					if(!domin[nn].SetEquals(D))
					{
						domin[nn] = D;
						change = true;
					}
				}
			}

			return domin;
		}
	}

	/// <summary>
	/// This class contains information on specifications for a single type.
	/// </summary>
	public sealed class NodeStructureSpec
	{
		/// <summary>
		/// The node that this spec is on.
		/// </summary>
		public readonly PIMNode OnNode;

		/// <summary>
		/// The set of exposed locations that are captured in this node (i.e. that this node dominates).
		/// </summary>
		public readonly HashSet<PIMNode> Captured;

		/// <summary>
		/// The set of locations that are exposed for the first time via this node.
		/// </summary>
		public readonly HashSet<Tuple<IRBasicAccessPath, PIMNode>> Exposed;

		/// <summary>
		/// The set of locations that are exposed by other types and the exposure is propagated by this node.
		/// </summary>
		public readonly HashSet<Tuple<IRBasicAccessPath, PIMNode>> Propagated;

		public NodeStructureSpec(PIMNode nn)
		{
			this.OnNode = nn;

			this.Captured = new HashSet<PIMNode>();

			this.Exposed = new HashSet<Tuple<IRBasicAccessPath, PIMNode>>();
			this.Propagated = new HashSet<Tuple<IRBasicAccessPath, PIMNode>>();
		}

		public void AddExposed(IRBasicAccessPath pspec, PIMNode tn)
		{
			this.Exposed.Add(new Tuple<IRBasicAccessPath, PIMNode>(pspec, tn));
		}

		public bool PropagateSpecsInto(NodeStructureSpec ospec, PIMEdge ee, PIModel pim, Dictionary<PIMNode, HashSet<PIMNode>> doms)
		{
			int psize = this.Propagated.Count;
			int csize = this.Captured.Count;
			
			HashSet<PIMNode> ddom = doms[this.OnNode];
			List<Tuple<IRBasicAccessPath, PIMNode>> npl = new List<Tuple<IRBasicAccessPath, PIMNode>>();
			npl.AddRange(ospec.Exposed);
			npl.AddRange(ospec.Propagated);
			npl.RemoveAll((Tuple<IRBasicAccessPath, PIMNode> rpl) => ospec.Captured.Contains(rpl.Item2));

			foreach(Tuple<IRBasicAccessPath, PIMNode> pn in npl)
			{
				if(DoDomCheck(this.OnNode, pn.Item2, doms))
					this.Captured.Add(pn.Item2);
				
				this.Propagated.Add(new Tuple<IRBasicAccessPath, PIMNode>(ee.OffsetInSource, pn.Item2));
			}

			return (this.Propagated.Count != psize) || (this.Captured.Count != csize);
		}

		private static bool DoDomCheck(PIMNode nn, PIMNode tn, Dictionary<PIMNode, HashSet<PIMNode>> doms)
		{
			if(!nn.LocationType.TypeDescription.IsBuiltinType && nn.IsUniqueLocation)
				return doms[nn].Contains(tn);
			else
			{
				return false;
			}
		}
	}
}
